import json
from typing import Callable, Any, Optional
import re
import requests
import jsonpath
from lxml import etree

type ExtractorFunc = Callable[[requests.Response, Optional[str]], Any]

EXTRACTORS: dict[str, ExtractorFunc] = {}


def extractor(name: str):
    def wrapper(func: ExtractorFunc):
        assert name not in EXTRACTORS
        EXTRACTORS[name] = func
        return func

    return wrapper


def extract_variables(response: requests.Response, extract_conf: dict[str, str]):
    my_vars = {}
    for var_name, rule in extract_conf.items():
        if ":" in rule:
            method, expr = rule.split(":", 1)
            method = method.strip()
            expr = expr.strip()
        else:
            method, expr = rule, None
            method = method.strip()

        assert method in EXTRACTORS
        my_vars[var_name] = EXTRACTORS[method](response, expr)
    return my_vars


@extractor("jsonpath")
def extract_jsonpath(response: requests.Response, expr: str):
    if expr is None:
        raise TypeError('jsonpath expr can not be None')
    matches = jsonpath.jsonpath(response.json(), expr)
    if not matches:
        return None
    if len(matches) == 1:
        return matches[0]
    return json.dumps(matches)


@extractor("xpath")
def extract_xpath(response: requests.Response, expr: str):
    html = etree.HTML(response.text)
    matches = html.xpath(expr)
    if not matches:
        return None
    if len(matches) == 1:
        return matches[0]
    return json.dumps(matches)


@extractor("regex")
def extract_regex(response: requests.Response, expr: str):
    matches = re.search(expr, response.text)
    return matches.group(0) if matches else None


@extractor("header")
def extract_header(response: requests.Response, expr: str):
    return response.headers.get(expr)


@extractor("status")
def extract_status(response: requests.Response, expr=None):
    return response.status_code
